"""
基于RSRS指标的螺纹钢策略
使用15分钟K线，单因子策略模板
"""

import numpy as np
import pandas as pd
from datetime import datetime
from tqsdk import TqApi, TargetPosTask
from sklearn.linear_model import LinearRegression

class OhDiffStrategy:
    strategy_name = "OhDiffNewStrategy"
    """
    螺纹钢策略
    
    开盘价收盘价差异的统计:
    1. 计算一段时间内开盘价和收盘价斜率
    2. 对斜率进行标准化处理
    3. 斜率越大表示上涨趋势越强，斜率越小表示下跌趋势越强
    
    策略逻辑:
    - oh_diff > 上阈值: 做多
    - oh_diff < 下阈值: 做空
    - 其他情况: 平仓或持有
    """
    
    def __init__(self, api, symbol="rb",exchange="SHFE", **kwargs):
        self.api = api
        self.symbol = symbol
        self.i_symbol = f"KQ.i@{exchange}.{symbol}"
        self.m_symbol = f"KQ.m@{exchange}.{symbol}"
        
        # RSRS策略参数
        self.corr_period = kwargs.get('corr_period', 30)  # RSRS计算周期
        self.corr_std_period = kwargs.get('corr_std_period', 600)  # 标准化周期
        self.long_threshold = kwargs.get('long_threshold', 0.7)  # 做多阈值
        self.short_threshold = kwargs.get('short_threshold', -0.7)  # 做空阈值
        self.close_threshold = kwargs.get('close_threshold', 0.2)  # 平仓阈值
        self.corr_series= []
        
        # 风控参数
        self.position_size = kwargs.get('position_size', 1)
        self.market_period = kwargs.get('market_period', 60)
        
        # 数据订阅 - 使用15分钟K线
        self.kline_15m = api.get_kline_serial(self.i_symbol, self.market_period * 60, data_length=800)
        
        # 交易对象
        self.main_symbol = api.get_quote(self.m_symbol).underlying_symbol
        self.target_pos = TargetPosTask(api, self.main_symbol)
        
        # 交易状态
        self.entry_price = 0.0
        self.stop_loss_price = 0.0
        self.take_profit_price = 0.0
        self.last_trade_time = 0
        
    
    def calculate_rsrs(self) -> float:
        """
        计算指标
        
        Returns:
            标准化后的指标值
        """
        try:
            if len(self.kline_15m) < self.corr_std_period:
                return 0.0
            # 获取最近的数据
            close_prices = self.kline_15m.close.values
            open_prices = self.kline_15m.open.values
            
            # 首次计算，从足够的历史数据开始计算
            if len(self.corr_series) < self.corr_std_period:
                for i in range(self.corr_period, len(close_prices)-1):
                    # 获取计算窗口内的高低价
                    window_close = close_prices[i-self.corr_period:i]
                    window_open = open_prices[i-self.corr_period:i]
                    mem = (window_close - window_open).mean()
                    corr = np.corrcoef(window_close, window_open)[0][1]*(1 if mem>=0 else -1)
                    self.corr_series.append(corr)
            # 后续更新计算
            else:
                    i = len(close_prices) - 1
                    window_close = close_prices[i-self.corr_period:i]
                    window_open = open_prices[i-self.corr_period:i]
                    mem = (window_close - window_open).mean()
                    corr = np.corrcoef(window_close, window_open)[0][1]*(1 if mem>=0 else -1)
                    self.corr_series.append(corr)

            current_corr = self.corr_series[-1]
            
            # 标准化处理
            recent_corrs = self.corr_series[-self.corr_std_period:]
            mean_corr = np.mean(recent_corrs)
            std_corr = np.std(recent_corrs)
            
            if std_corr > 0:
                signal_score = (current_corr - mean_corr) / std_corr
            else:
                signal_score = 0.0
            
            return signal_score
            
        except Exception as e:
            print(f"计算失败: {e}")
            return 0.0
    
    def calculate_atr(self, period: int = 14) -> float:
        """计算ATR用于止损止盈"""
        try:
            if len(self.kline_15m) < period + 1:
                return self.kline_15m.close.iloc[-1] * 0.02  # 默认2%
            
            high = self.kline_15m.high.iloc[-period:]
            low = self.kline_15m.low.iloc[-period:]
            close = self.kline_15m.close.iloc[-period-1:-1]  # 前一日收盘价
            
            tr1 = high - low
            tr2 = abs(high - close.values)
            tr3 = abs(low - close.values)
            
            tr = np.maximum(tr1, np.maximum(tr2, tr3))
            atr = np.mean(tr)
            
            return atr
            
        except Exception as e:
            return self.kline_15m.close.iloc[-1] * 0.02
    
    def on_bar(self):
        """主循环"""
        bar_count = 0
        while True:
            self.api.wait_update()
            if self.api.is_changing(self.kline_15m):
                bar_count += 1
                current_time = self.kline_15m.datetime.iloc[-1]
                current_price = self.kline_15m.close.iloc[-1]
                volume = self.kline_15m.volume.iloc[-1]
                
                # 每50根K线打印一次进度
                if bar_count % 50 == 0:
                    print(f"处理第{bar_count}根15分钟K线: 时间={current_time}, 价格={current_price:.1f}")
                
                if volume == 0:
                    self.run()
    
    def run(self):
        """策略执行逻辑"""
        try:
            # 切换主力合约
            self.main_symbol = self.api.get_quote(self.m_symbol).underlying_symbol
            self.target_pos = TargetPosTask(self.api, self.main_symbol)
            
            # 平掉非当前主力的持仓
            positions = self.api.get_position()
            for sym, pos in positions.items():
                if sym != self.main_symbol and pos.pos != 0:
                    TargetPosTask(self.api, sym).set_target_volume(0)
            
            main_pos = positions.get(self.main_symbol, None)
            current_pos = main_pos.pos if main_pos else 0
            current_price = self.kline_15m.close.iloc[-1]
            current_time = self.kline_15m.datetime.iloc[-1]
            
            # 检查数据是否足够
            if len(self.kline_15m) < self.corr_std_period:
                return
            
            # 控制交易频率 (15分钟间隔)
            if current_time - self.last_trade_time < 15 * 60 * 1000:
                return
            
            # 计算指标
            signal_score = self.calculate_rsrs()

            # 计算ATR用于止损止盈
            atr_val = self.calculate_atr()
            
            # 调试信息：每20次运行打印一次
            if hasattr(self, 'run_count'):
                self.run_count += 1
            else:
                self.run_count = 1
                
            if self.run_count % 20 == 0:
                print(f"{self.strategy_name}策略运行{self.run_count}次: 指标值={signal_score:.3f}, 持仓={current_pos}, 价格={current_price:.1f}")
            
            if current_pos == 0:  # 无持仓
                if signal_score > self.long_threshold:  # 做多信号
                    # 做多
                    self.entry_price = current_price
                    self.target_pos.set_target_volume(self.position_size)
                    self.last_trade_time = current_time
                    
                    print(f"🟢 做多: 指标={signal_score:.3f}, 价格={current_price:.1f}")
                
                elif signal_score < self.short_threshold:  # 做空信号
                    # 做空
                    self.entry_price = current_price
                    self.target_pos.set_target_volume(-self.position_size)
                    self.last_trade_time = current_time
                    
                    print(f"🔴 做空: 指标={signal_score:.3f}, 价格={current_price:.1f}")
            
            else:  # 有持仓
                should_close = False
                close_reason = ""
                
                if current_pos > 0:  # 多头持仓
                    if signal_score < -self.close_threshold:  # 达到离场信号
                        should_close = True
                        close_reason = f"离场信号({signal_score:.3f})"
                
                elif current_pos < 0:  # 空头持仓
                    if signal_score > self.close_threshold:  # 达到离场信号
                        should_close = True
                        close_reason = f"离场信号({signal_score:.3f})"
                
                if should_close:
                    # 平仓
                    self.target_pos.set_target_volume(0)
                    self.last_trade_time = current_time
                    
                    print(f"⚪ {close_reason}: 价格={current_price:.1f}")
                    
        except Exception as e:
            print(f"策略执行失败: {e}")
    